# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: donggu

from experimental.prophet.ops import timeseries
from experimental.prophet.ops.control_flow import cond
from experimental.prophet.ops.constant import constant
from experimental.prophet.ops.elemwise_math import cbrt
from experimental.prophet.contrib.variable import last_valid


def vwap(trade, window_size):
  moving_product = timeseries.time_moving_window(trade.price * trade.qty, window_size).sum()
  moving_qty = timeseries.time_moving_window(trade.qty, window_size).sum()
  # vwap = cond(avg_qty > 0, avg_product / avg_qty, constant(np.nan))
  return moving_product / moving_qty


def vwap_deviation(trade, window_size):
  vwap_ = vwap(trade, window_size)
  last_price = last_valid(trade.price)
  return last_valid(cbrt((vwap_ - last_price) / vwap_))


__all__ = ['vwap', 'vwap_deviation']
